"""
新的API路由 - 支持多表关系结构 - MySQL版本
"""

from aiohttp import web
import pymysql
import os
import json
from urllib.parse import unquote

class DatabaseManager:
    """MySQL数据库管理器"""
    
    def __init__(self):
        self.mysql_config = {
            'host': 'localhost',
            'port': 3306,
            'user': 'root',
            'password': 'root',
            'database': 'comfyui-xhh',
            'charset': 'utf8mb4'
        }
    
    def get_connection(self):
        """获取MySQL数据库连接"""
        try:
            return pymysql.connect(**self.mysql_config)
        except Exception as e:
            print(f"MySQL连接失败: {e}")
            return None

class ImageAPI:
    """图片相关API"""
    
    def __init__(self, db_manager):
        self.db = db_manager
    
    async def get_image_info(self, request):
        """获取图片信息"""
        try:
            image_id = request.match_info['image_id']
            
            conn = self.db.get_connection()
            if not conn:
                return web.json_response({
                    'success': False,
                    'message': '数据库连接失败'
                })
            
            try:
                with conn.cursor() as cursor:
                    cursor.execute("""
                        SELECT id, filename, file_path, file_size, width, height, mime_type, created_at
                        FROM images WHERE id = %s
                    """, (image_id,))
                    
                    image = cursor.fetchone()
                    if not image:
                        return web.json_response({
                            'success': False,
                            'message': '图片不存在'
                        })
                    
                    return web.json_response({
                        'success': True,
                        'data': {
                            'id': image[0],
                            'filename': image[1],
                            'file_path': image[2],
                            'file_size': image[3],
                            'width': image[4],
                            'height': image[5],
                            'mime_type': image[6],
                            'created_at': image[7]
                        }
                    })
            finally:
                conn.close()
                
        except Exception as e:
            return web.json_response({
                'success': False,
                'message': str(e)
            })
    
    async def get_image_loras(self, request):
        """获取图片关联的LoRA列表"""
        try:
            image_id = request.match_info['image_id']
            
            conn = self.db.get_connection()
            if not conn:
                return web.json_response({
                    'success': False,
                    'message': '数据库连接失败'
                })
            
            try:
                with conn.cursor() as cursor:
                    cursor.execute("""
                        SELECT l.id, l.name, l.path, l.size, l.tags, l.description, l.trigger, l.md5, l.is_favorite
                        FROM loras l
                        JOIN lora_images li ON l.id = li.lora_id
                        WHERE li.image_id = %s
                        ORDER BY li.created_at DESC
                    """, (image_id,))
                    
                    loras = []
                    for row in cursor.fetchall():
                        loras.append({
                            'id': row[0],
                            'name': row[1],
                            'path': row[2],
                            'size': row[3],
                            'tags': json.loads(row[4]) if row[4] else [],
                            'description': row[5],
                            'trigger': row[6],
                            'prompt': row[7],
                            'is_favorite': bool(row[8])
                        })
                    
                    return web.json_response({
                        'success': True,
                        'data': loras
                    })
            finally:
                conn.close()
                
        except Exception as e:
            return web.json_response({
                'success': False,
                'message': str(e)
            })

class LoraAPI:
    """LoRA相关API"""
    
    def __init__(self, db_manager):
        self.db = db_manager
    
    async def get_lora_images(self, request):
        """获取LoRA关联的图片列表"""
        try:
            lora_id = request.match_info['lora_id']
            
            conn = self.db.get_connection()
            if not conn:
                return web.json_response({
                    'success': False,
                    'message': '数据库连接失败'
                })
            
            try:
                with conn.cursor() as cursor:
                    cursor.execute("""
                        SELECT i.id, i.filename, i.file_path, i.file_size, i.width, i.height, i.mime_type
                        FROM images i
                        JOIN lora_images li ON i.id = li.image_id
                        WHERE li.lora_id = %s
                        ORDER BY li.created_at DESC
                    """, (lora_id,))
                    
                    images = []
                    for row in cursor.fetchall():
                        images.append({
                            'id': row[0],
                            'filename': row[1],
                            'file_path': row[2],
                            'file_size': row[3],
                            'width': row[4],
                            'height': row[5],
                            'mime_type': row[6]
                        })
                    
                    return web.json_response({
                        'success': True,
                        'data': images
                    })
            finally:
                conn.close()
                
        except Exception as e:
            return web.json_response({
                'success': False,
                'message': str(e)
            })
    
    async def get_lora_images_by_name(self, request):
        """根据LoRA名称获取图片列表"""
        try:
            lora_name = request.query.get('name', '')
            if not lora_name:
                return web.json_response({
                    'success': False,
                    'message': '缺少LoRA名称参数'
                })
            
            conn = self.db.get_connection()
            if not conn:
                return web.json_response({
                    'success': False,
                    'message': '数据库连接失败'
                })
            
            try:
                with conn.cursor() as cursor:
                    cursor.execute("""
                        SELECT i.id, i.filename, i.file_path, i.file_size, i.width, i.height, i.mime_type
                        FROM images i
                        JOIN lora_images li ON i.id = li.image_id
                        JOIN loras l ON li.lora_id = l.id
                        WHERE l.name LIKE %s
                        ORDER BY li.created_at DESC
                    """, (f'%{lora_name}%',))
                    
                    images = []
                    for row in cursor.fetchall():
                        images.append({
                            'id': row[0],
                            'filename': row[1],
                            'file_path': row[2],
                            'file_size': row[3],
                            'width': row[4],
                            'height': row[5],
                            'mime_type': row[6]
                        })
                    
                    return web.json_response({
                        'success': True,
                        'data': images
                    })
            finally:
                conn.close()
                
        except Exception as e:
            return web.json_response({
                'success': False,
                'message': str(e)
            })

class PromptAPI:
    """提示词相关API"""
    
    def __init__(self, db_manager):
        self.db = db_manager
    
    async def get_prompt_images(self, request):
        """获取提示词的图片列表"""
        try:
            prompt_id = request.match_info['prompt_id']
            
            conn = self.db.get_connection()
            if not conn:
                return web.json_response({
                    'success': False,
                    'message': '数据库连接失败'
                })
            
            try:
                with conn.cursor() as cursor:
                    cursor.execute("""
                        SELECT i.id, i.filename, i.file_path, i.file_size, i.width, i.height, i.mime_type
                        FROM images i
                        JOIN prompt_images pi ON i.id = pi.image_id
                        WHERE pi.prompt_id = %s
                        ORDER BY pi.created_at DESC
                    """, (prompt_id,))
                    
                    images = []
                    for row in cursor.fetchall():
                        images.append({
                            'id': row[0],
                            'filename': row[1],
                            'file_path': row[2],
                            'file_size': row[3],
                            'width': row[4],
                            'height': row[5],
                            'mime_type': row[6]
                        })
                    
                    return web.json_response({
                        'success': True,
                        'data': images,
                        'images': images  # 保持向后兼容
                    })
            finally:
                conn.close()
                
        except Exception as e:
            return web.json_response({
                'success': False,
                'message': str(e)
            })
    
    async def get_image_prompt_info(self, request):
        """根据图片获取提示词信息"""
        try:
            image_url = request.query.get('image', '')
            if not image_url:
                return web.json_response({
                    'success': False,
                    'message': '缺少图片参数'
                })
            
            conn = self.db.get_connection()
            if not conn:
                return web.json_response({
                    'success': False,
                    'message': '数据库连接失败'
                })
            
            try:
                with conn.cursor() as cursor:
                    cursor.execute("""
                        SELECT p.id, p.text, p.source, p.created_at
                        FROM prompts p
                        JOIN prompt_images pi ON p.id = pi.prompt_id
                        JOIN images i ON pi.image_id = i.id
                        WHERE i.file_path = %s
                        ORDER BY pi.created_at DESC
                        LIMIT 1
                    """, (image_url,))
                    
                    result = cursor.fetchone()
                    if result:
                        return web.json_response({
                            'success': True,
                            'data': {
                                'id': result[0],
                                'text': result[1],
                                'source': result[2],
                                'created_at': result[3]
                            }
                        })
                    else:
                        return web.json_response({
                            'success': False,
                            'message': '未找到相关提示词'
                        })
            finally:
                conn.close()
                
        except Exception as e:
            return web.json_response({
                'success': False,
                'message': str(e)
            })

class RelationshipAPI:
    """关系管理API"""
    
    def __init__(self, db_manager):
        self.db = db_manager
    
    async def add_lora_image(self, request):
        """添加LoRA图片关联"""
        try:
            data = await request.json()
            lora_id = data.get('lora_id')
            image_id = data.get('image_id')
            
            if not lora_id or not image_id:
                return web.json_response({
                    'success': False,
                    'message': '缺少必要参数'
                })
            
            conn = self.db.get_connection()
            if not conn:
                return web.json_response({
                    'success': False,
                    'message': '数据库连接失败'
                })
            
            try:
                with conn.cursor() as cursor:
                    # 检查关联是否已存在
                    cursor.execute("""
                        SELECT id FROM lora_images 
                        WHERE lora_id = %s AND image_id = %s
                    """, (lora_id, image_id))
                    
                    if cursor.fetchone():
                        return web.json_response({
                            'success': False,
                            'message': '关联已存在'
                        })
                    
                    # 创建新关联
                    cursor.execute("""
                        INSERT INTO lora_images (lora_id, image_id, created_at)
                        VALUES (%s, %s, NOW())
                    """, (lora_id, image_id))
                    
                    conn.commit()
                    
                    return web.json_response({
                        'success': True,
                        'message': '关联创建成功'
                    })
            finally:
                conn.close()
                
        except Exception as e:
            return web.json_response({
                'success': False,
                'message': str(e)
            })
    
    async def add_prompt_image(self, request):
        """添加提示词图片关联"""
        try:
            data = await request.json()
            prompt_id = data.get('prompt_id')
            image_id = data.get('image_id')
            
            if not prompt_id or not image_id:
                return web.json_response({
                    'success': False,
                    'message': '缺少必要参数'
                })
            
            conn = self.db.get_connection()
            if not conn:
                return web.json_response({
                    'success': False,
                    'message': '数据库连接失败'
                })
            
            try:
                with conn.cursor() as cursor:
                    # 检查关联是否已存在
                    cursor.execute("""
                        SELECT id FROM prompt_images 
                        WHERE prompt_id = %s AND image_id = %s
                    """, (prompt_id, image_id))
                    
                    if cursor.fetchone():
                        return web.json_response({
                            'success': False,
                            'message': '关联已存在'
                        })
                    
                    # 创建新关联
                    cursor.execute("""
                        INSERT INTO prompt_images (prompt_id, image_id, created_at)
                        VALUES (%s, %s, NOW())
                    """, (prompt_id, image_id))
                    
                    conn.commit()
                    
                    return web.json_response({
                        'success': True,
                        'message': '关联创建成功'
                    })
            finally:
                conn.close()
                
        except Exception as e:
            return web.json_response({
                'success': False,
                'message': str(e)
            })

def setup_routes(app):
    """设置API路由"""
    # 创建数据库管理器
    db_manager = DatabaseManager()
    
    # 创建API实例
    image_api = ImageAPI(db_manager)
    lora_api = LoraAPI(db_manager)
    prompt_api = PromptAPI(db_manager)
    relationship_api = RelationshipAPI(db_manager)
    
    # 图片相关路由
    app.router.add_get('/api/xhh/images/{image_id}', image_api.get_image_info)
    app.router.add_get('/api/xhh/images/{image_id}/loras', image_api.get_image_loras)
    
    # LoRA相关路由
    app.router.add_get('/api/xhh/loras/{lora_id}/images', lora_api.get_lora_images)
    app.router.add_get('/api/xhh/lora-images', lora_api.get_lora_images_by_name)
    
    # 提示词相关路由
    app.router.add_get('/api/xhh/prompts/{prompt_id}/images', prompt_api.get_prompt_images)
    app.router.add_get('/api/xhh/image-prompt-info', prompt_api.get_image_prompt_info)
    
    # 关系管理路由
    app.router.add_post('/api/xhh/lora-images', relationship_api.add_lora_image)
    app.router.add_post('/api/xhh/prompt-images', relationship_api.add_prompt_image)
    
    print("✅ MySQL API路由注册完成")
